package cn.apelx.shardingjdbc.config;

import cn.hutool.core.collection.CollectionUtil;
import org.apache.shardingsphere.infra.config.RuleConfiguration;
import org.apache.shardingsphere.sharding.algorithm.config.AlgorithmProvidedShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingAutoTableRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;

import java.util.*;
import java.util.stream.Collectors;

/**
 * Sharding后置处理RuleConfig
 *
 * @author apelx
 * @since 2022-05-14
 */
public class ShardingPostProcess {

    private final String defaultDatasourceName;

    ShardingPostProcess(String defaultDatasourceName) {
        this.defaultDatasourceName = defaultDatasourceName;
    }

    /**
     * 对所有分片规则中的表与当前系统中所有实体类表进行对比
     * 没有配置规则在shardingRule中, 添加默认规则,使用指定的默认数据源
     *
     * @param rules
     */
    public void postProcess(RuleConfiguration rules) {
        if (rules instanceof AlgorithmProvidedShardingRuleConfiguration) {

            Collection<ShardingAutoTableRuleConfiguration> autoTables = Optional.ofNullable(((AlgorithmProvidedShardingRuleConfiguration) rules).getAutoTables()).orElse(new ArrayList<>(0));
            Collection<String> broadcastTables = Optional.ofNullable(((AlgorithmProvidedShardingRuleConfiguration) rules).getBroadcastTables()).orElse(new ArrayList<>(0));
            Collection<ShardingTableRuleConfiguration> tables = Optional.ofNullable(((AlgorithmProvidedShardingRuleConfiguration) rules).getTables()).orElse(new ArrayList<>(0));

            // 已配置分片规则的表
            Set<String> shardingTables = autoTables.stream().map(ShardingAutoTableRuleConfiguration::getLogicTable).collect(Collectors.toSet());
            shardingTables.addAll(new HashSet<>(broadcastTables));
            shardingTables.addAll(tables.stream().map(ShardingTableRuleConfiguration::getLogicTable).collect(Collectors.toSet()));

            // 系统所有实体表名
            Set<String> allTables = EntityBeanUtil.getAllTableNames();

            // 取差集
            Collection<String> differenceSet = CollectionUtil.subtract(allTables, shardingTables);
            System.err.println("自动设置默认数据源tables[" + differenceSet + "]");
            if (CollectionUtil.isEmpty(differenceSet)) {
                return;
            }

            if (((AlgorithmProvidedShardingRuleConfiguration) rules).getTables() == null) {
                ((AlgorithmProvidedShardingRuleConfiguration) rules).setTables(new ArrayList<>(differenceSet.size()));
            }
            for (String tableName : differenceSet) {
                ShardingTableRuleConfiguration config =
                        new ShardingTableRuleConfiguration(tableName, defaultDatasourceName + "." + tableName);
                ((AlgorithmProvidedShardingRuleConfiguration) rules).getTables().add(config);
            }
        }
    }
}